"""
《邢不行-2023新版|Python股票量化投资课程》
author: 邢不行
微信: xbx9585

选股使用的过滤的脚本
"""

from Config import *


def filter_and_rank(df, stg_name):
    # print(df['营业收入同比增速_60'].tail(6))
    # exit()
    """
    过滤股票 & 选股票
    :param df: 原始数据
    :param stg_name: 策略名称，根据名称选择需要运行的策略
    :return: 返回选股结果
    """
    # =====针对不同的策略进行选股
    #优化方向0：原始策略
    if stg_name == '小市值_原始策略':
        # 计算总市值排名
        df['总市值排名'] = df.groupby('交易日期')['总市值'].rank(ascending=True, method='min')
        # 计算复合因子
        df['复合因子'] = df['总市值排名']

    # 优化方向1：用量价因子对小市值进行优化
    if stg_name == '小市值_量价优化':
        # 计算总市值排名
        df['总市值排名'] = df.groupby('交易日期')['总市值'].rank(ascending=True, method='min')
        # 计算alpha95排名  衡量低波动
        df['alpha95排名'] = df.groupby('交易日期')['alpha95'].rank(ascending=True, method='min')
        # 计算Ret20排名  衡量超跌
        df['Ret_20排名'] = df.groupby('交易日期')['Ret_20'].rank(ascending=True, method='min')
        # df['均线趋势排名'] = df.groupby('交易日期')['均线趋势'].rank(ascending=False, method='min')

        # df['涨跌幅'] = df['收盘价'] / df['前收盘价'] - 1
        # df.loc[(df['均线_5'] - df['均线_5'].shift(1)) > 0, '均线_5趋势'] = 1
        # df.loc[(df['均线_5'] - df['均线_5'].shift(1)) < 0, '均线_5趋势'] = -1
        # df['均线_5趋势'].ffill(inplace=True)
        # df.loc[(df['均线_10'] - df['均线_10'].shift(1)) > 0, '均线_10趋势'] = 1
        # df.loc[(df['均线_10'] - df['均线_10'].shift(1)) < 0, '均线_10趋势'] = -1
        # df['均线_10趋势'].ffill(inplace=True)
        # df.loc[(df['均线_20'] - df['均线_20'].shift(1)) > 0, '均线_20趋势'] = 1
        # df.loc[(df['均线_20'] - df['均线_20'].shift(1)) < 0, '均线_20趋势'] = -1
        # df['均线_20趋势'].ffill(inplace=True)
        # df.loc[(df['均线_30'] - df['均线_30'].shift(1)) > 0, '均线_30趋势'] = 1
        # df.loc[(df['均线_30'] - df['均线_30'].shift(1)) < 0, '均线_30趋势'] = -1
        # df['均线_30趋势'].ffill(inplace=True)
        # df.loc[(df['均线_60'] - df['均线_60'].shift(1)) > 0, '均线_60趋势'] = 1
        # df.loc[(df['均线_60'] - df['均线_60'].shift(1)) < 0, '均线_60趋势'] = -1
        # df['均线_60趋势'].ffill(inplace=True)
        # df.loc[(df['均线_88'] - df['均线_88'].shift(1)) > 0, '均线_88趋势'] = 1
        # df.loc[(df['均线_88'] - df['均线_88'].shift(1)) < 0, '均线_88趋势'] = -1
        # df['均线_88趋势'].ffill(inplace=True)
        # df['均线趋势'] = df['均线_5趋势'] + df['均线_10趋势'] + df['均线_20趋势'] + df['均线_30趋势'] + df['均线_60'] + df['均线_88趋势']
        # df['均线趋势排名'] = df['均线趋势'].rank(ascending=True, method='min')

         # 计算5、10、20日均线间的最大差值
        # df['三线乖离率'] = (df[['均线_5', '均线_10', '均线_20']].T.max() / df[
        #     ['均线_5', '均线_10', '均线_20']].T.min() - 1) * 100
        # 计算5、10、20、30日均线间的最大差值
        # df['四线乖离率'] = (df[['均线_5', '均线_10', '均线_20', '30日均线']].T.max() / df[
        #     ['均线_5', '均线_10', '均线_20', '30日均线']].T.min() - 1) * 100


        # df['涨停价'] = df['前收盘价'] * 1.1
        # df['跌停价'] = df['前收盘价'] * 0.9

        # df['收盘涨停'] = 0
        # df.loc[df['收盘价'] >= df['涨停价'], '收盘涨停'] = 1
        # df['收盘跌停'] = 0
        # df.loc[df['收盘价'] <= df['跌停价'], '收盘跌停'] = 1
        # c0_1 = df['收盘跌停'].rolling(9).sum() < 1
        # c0_2 = df['收盘涨停'].rolling(19).sum() <= 2

        # c1_0 = (df['最高价'] / df['收盘价'] > 1.09) & (df['涨跌幅'] < 1)
        # df['冲高回落'] = 0
        # df.loc[c1_0, '冲高回落'] = 1
        # c1_1 = df['冲高回落'].rolling(9).sum() < 1
        # condition_base = c0_1 & c0_2 & c1_1

        # 回踩日均线
        # c_20 = (df['涨跌幅'] < 5) & (df['涨跌幅'] > -5) & (df['收盘价'] > df['开盘价']) & (
        #         df['最高价'] > df['均线_88']) & (df['收盘价'] / df['均线_88'] < 1.05) & (
        #             df['最低价'] <= df['均线_88']) & (df['均线_88趋势'] == 1) & (df['均线_88'] < df['均线_60']) & (df['均线_60'] < df['均线_30'])
        # c_20 = (df['收盘价'] > df['开盘价']) & (df['最高价'] > df['均线_88']) & (df['收盘价'] / df['均线_88'] < 1.05) & (
        #             df['最低价'] <= df['均线_88']) & (df['均线_88趋势'] == 1) & (df['均线_88'] < df['均线_60']) & (df['均线_60'] < df['均线_30'])
        # df['回踩ma'] = 0
        # df.loc[c_20, '回踩ma'] = 1

        # df['上穿三线'] = 0
        # cross_3_1 = (df['三线乖离率'] <= 1.8) & (df['最高价'] >= df['均线_88']) & (
        #         df['最高价'] >= df['均线_60']) & (df['最高价'] >= df['均线_30']) & (df['涨跌幅'] > -2) & (
        #                     df['收盘价'] > df['开盘价']) & (df['最低价'] <= df['均线_88']) & (
        #                     df['最低价'] <= df['均线_60']) & (
        #                     df['最低价'] <= df['均线_30'])
        # cross_3_2 = (df['均线_60趋势'] == 1)
        # df.loc[c1_0 & cross_3_1 & cross_3_2, '上穿三线'] = 1

        # df = df[(df['三线乖离率'] <= 1.8)]

        # df = df[df['均线_20趋势'] == 1]
        # df = df[(df['三线乖离率'] <= 1.8) & (df['收盘价'] >= df['均线_20'])]
        

        # 计算复合因子
        df['复合因子'] = df['alpha95排名'] + df['总市值排名'] + df['Ret_20排名']

    # 优化方向2：用基本面数据对小市值进行优化
    if stg_name == '小市值_基本面优化':

        # 计算ROE百分比排名，去除ROE较差的20%的股票
        df['ROE排名'] = df.groupby('交易日期')['ROE'].rank(ascending=False, method='min', pct=True)
        df = df[df['ROE排名'] < 0.8]
        df = df[df['R_np_atoopc@xbx'] > 0]
        # df = df[df['R_total_compre_income_atsopc@xbx'] > 0]
        # df = df[df['R_basic_eps@xbx'] > 0]
        # df = df[df['净利润(当季)'] > 0]
        # 计算总市值排名
        df['总市值排名'] = df.groupby('交易日期')['总市值'].rank(ascending=True, method='min')
        # 计算归母净利润同比增速排名
        df['归母净利润同比增速排名'] = df.groupby('交易日期')['归母净利润同比增速_60'].rank(ascending=False, method='min')


        # 每股受益同比增速排名
        # df['每股受益同比增速排名'] = df.groupby('交易日期')['每股受益同比增速_60'].rank(ascending=False, method='min')
        # 营业收入同比增速排名
        # df = df[df['R_revenue@xbx'] > 0]
        # df['营业收入同比增速排名'] = df.groupby('交易日期')['营业收入同比增速_60'].rank(ascending=False, method='min')
        # df['均线趋势排名'] = df.groupby('交易日期')['均线趋势'].rank(ascending=True, method='min')
        

        # df = df[df['近期有试盘行为'] > 0]

        # 奈绪选股
        # ascending = True
        # df = df[(df['J'] > 80) & (df['D'] > 80)]
        # df['L97_排名'] = df.groupby('交易日期')['L97'].rank(ascending=ascending)
        # df['成交额std_10_排名'] = df.groupby('交易日期')['成交额std_10'].rank()
        # df['成交额std_20_排名'] = df.groupby('交易日期')['成交额std_20'].rank(ascending=ascending)
        # df['量价相关系数_10_排名'] = df.groupby('交易日期')['量价相关性'].rank(pct=True)
        # # df['量稳换手率变化率_排名'] = df.groupby('交易日期')['量稳换手率变化率'].rank(ascending=False)
        # df['异常换手率_排名'] = df.groupby('交易日期')['异常换手率'].rank(ascending=ascending)
        # df['bias_5_排名'] = df.groupby('交易日期')['bias_5'].rank()
        # df['bias_10_排名'] = df.groupby('交易日期')['bias_10'].rank()
        # df['MACD_Hist'] = df.groupby('交易日期')['MACD_Hist'].rank(pct=True)
        # df = df[df['MACD_Hist'] > 0.08]
        # # df['因子'] = 0.2 * df['成交额std_10_排名'] + 0.8 * df['bias_5_排名'] + 0.6 * df['L97_排名']\
        # #         + 1.8 * df['成交额std_20_排名'] + df['量稳换手率变化率_排名'] + 0.7 * df['bias_10_排名']
        # df['因子'] = 0.2 * df['成交额std_10_排名'] + 0.8 * df['bias_5_排名'] + 0.6 * df['L97_排名']\
        #         + 1.8 * df['成交额std_20_排名'] + 0.7 * df['bias_10_排名']
        
        # df['奈绪排名'] = df.groupby('交易日期')['因子'].rank()


        # df = df[(df['大户资金净流入'] > 0) & (df['big_cash_ma5'] > 0) & (df['big_cash_ma10'] > 0) &
        # (df['mid_cash_ma10'] < 0) & (df['mid_cash_ma5'] > 0)]
        # df = df[df['上影线'] < 0.05]

         # 计算5、10、20日均线间的最大差值
        # df['三线乖离率'] = (df[['均线_5', '均线_10', '均线_20']].T.max() / df[
        #     ['均线_5', '均线_10', '均线_20']].T.min() - 1) * 100
        # 计算5、10、20、30日均线间的最大差值
        # df['四线乖离率'] = (df[['均线_5', '均线_10', '均线_20', '30日均线']].T.max() / df[
        #     ['均线_5', '均线_10', '均线_20', '30日均线']].T.min() - 1) * 100

        # df = df[(df['三线乖离率'] <= 1.8)]

        # 计算复合因子
        df['复合因子'] = df['总市值排名'] + df['归母净利润同比增速排名'] 

    # 优化方向3：对小市值进行过滤
    if stg_name == '小市值_过滤优化':
        # 计算ROE百分比排名
        df['ROE排名'] = df.groupby('交易日期')['ROE'].rank(ascending=False, method='min', pct=True)
        # 计算归母净利润同比增速百分比排名
        df['归母净利润同比增速_60排名'] = df.groupby('交易日期')['归母净利润同比增速_60'].rank(ascending=False, method='min',
                                                                       pct=True)
        # 去除ROE较差的20%的股票
        df = df[df['ROE排名'] < 0.8]
        # 保留计算归母净利润同比增速排名靠前20%的股票
        df = df[df['归母净利润同比增速_60排名'] < 0.2]
        # 计算总市值排名
        df['总市值排名'] = df.groupby('交易日期')['总市值'].rank(ascending=True, method='min')
        df = df[df['R_np_atoopc@xbx'] > 0]
        # 计算复合因子
        df['复合因子'] = df['总市值排名']

    # 优化方向4：在限定股票池内选股
    if stg_name == '小市值_限定股票池':
        # 指数名称，可以替换的
        index_name = '中证1000'  # 沪深300   中证1000
        # 只保留成分股数据
        df = df[df[f'{index_name}成分股'] == 'Y']
        # 计算总市值排名
        df['总市值排名'] = df.groupby('交易日期')['总市值'].rank(ascending=True, method='min')
        # 计算复合因子
        df['复合因子'] = df['总市值排名']

    # 优化方向5：用量价因子对小市值进行优化突破
    if stg_name == '小市值_量价优化突破':
        # 计算总市值排名
        df['总市值排名'] = df.groupby('交易日期')['总市值'].rank(ascending=True, method='min')
        # 计算alpha95排名  衡量低波动
        df['alpha95排名'] = df.groupby('交易日期')['alpha95'].rank(ascending=True, method='min')
        # 计算Ret20排名  衡量超跌
        df['Ret_20排名'] = df.groupby('交易日期')['Ret_20'].rank(ascending=True, method='min')

        df['closeShift1'] = df['收盘价'].shift(1)
        df['closeShift2'] = df ['收盘价'].shift(2)
        df['closeShift3'] = df['收盘价'].shift(3)
        df['closeShift4']= df['收盘价'].shift(4)
        df['upperShift1'] = df['upper'].shift(1)
        df['upperShift2'] = df['upper'].shift(2)
        df['upperShift3']= df['upper'].shift(3)
        df['upperShift4'] = df['upper'].shift(4)
        df['midShift1'] = df['median'].shift(1)
        df['midShift2'] = df['median'].shift(2)
        df['midShift3'] = df['median'].shift(3)
        df['midShift4'] = df['median'] .shift(4)
        df['volMa5Sft1'] = df['vol_m5'].shift(1)
        df['volMa65Sft1'] = df['vol_m65'].shift(1)
        df =  df[df['closeShift1'] > df['midShift1']]
        df = df[df['closeShift2'] > df['midShift2']]
        df = df[df['closeShift3'] > df['midShift3']]
        df = df[df['closeShift4'] > df['midShift4']]
        df = df[df['closeShift1'] < df['upperShift1']]
        df = df[df['closeShift2'] < df['upperShift2']]
        df = df[df['closeShift3'] < df['upperShift3']]
        df = df[df['closeShift4'] < df['upperShift4']]
        df = df[df['收盘价'] > df['upper']]
        #df = df[df[volMa5Sf1] <= df[volMa65Sf1]]#前-天MA5小于MA65df= dfldf[vol ma5]>= df['vol ma65]]#当天MA5大于MA65(金又)#df = dfldf[vol chg]>0]# 成交量涨幅大于0%# df-df[df[vol chg]<=5] # 成交量涨幅小于等于 10%
        df= df[df['vol_m5'] >= df['vol_m65']]
        df= df[df['开盘价']>df['median']]#开盘价大于MA20
        df = df[df['收盘价']> df['开盘价']] # 当日上涨
        df = df[df['收盘价']>df['前收盘价']]#当日真阳线
        df = df[df['MA20斜率']> 0] # MA20角度向上

        # 计算复合因子
        df['复合因子'] = df['alpha95排名'] + df['总市值排名'] + df['Ret_20排名']

    # 优化方向6：财务优化&低价股尾部去除
    if stg_name == '小市值_基本面优化低价尾部去除':
         # 计算ROE百分比排名，去除ROE较差的20%的股票
        df['ROE排名'] = df.groupby('交易日期')['ROE'].rank(ascending=False, method='min', pct=True)
        df = df[df['R_np_atoopc@xbx'] > 0]
        # 计算总市值排名
        df['总市值排名'] = df.groupby('交易日期')['总市值'].rank(ascending=True, method='min')
        # 计算归母净利润同比增速排名
        df['归母净利润同比增速排名'] = df.groupby('交易日期')['归母净利润同比增速_60'].rank(ascending=False, method='min',pct=True)


        #低价股尾部去除-------------------------------------------
        df = df[df['总市值'] < 3e9]
        df = df[df['换手率_5'] > 0.01]
        condition = (df['归母净利润同比增速排名'] < 0.8)
        df['收盘价_排名'] = df.groupby('交易日期')['收盘价'].rank(ascending=False, method='min', pct=True)
        condition &= (df['收盘价_排名'] < 0.95)
        df['ROE排名'] = df.groupby('交易日期')['ROE'].rank(ascending=False, method='min', pct=True)
        condition &= (df['ROE排名'] < 0.95)
        # df["WR_5_排名"] = df.groupby("交易日期")["WR_5"].rank(ascending=True, method="min", pct=True)
        # condition &= (df["WR_5_排名"] > 0.2)
        # df["JS_5_排名"] = df.groupby("交易日期")["JS_5"].rank(ascending=True, method="min", pct=True)
        # condition &= (df["JS_5_排名"] < 0.8)
        # 计算总市值排名
        df['总市值排名'] = df.groupby('交易日期')['总市值'].rank(ascending=True, method='min')
        # 计算Ret20排名  衡量超跌
        df['Ret_21排名'] = df.groupby('交易日期')['Ret_21'].rank(ascending=True, method='min')
        # 计算成交额
        df['成交额std_20_排名'] = df.groupby('交易日期')['成交额std_20'].rank(ascending=True, method='min')
        # 中户资金
        df['中户买入占比排名'] = df.groupby('交易日期')['中户买入占比'].rank(ascending=False, method='min')
        #低价股尾部去除-------------------------------------------

        df['复合因子'] = df['总市值排名'] + df['中户买入占比排名'] + df['成交额std_20_排名'] + df['Ret_21排名']

    # 优化方向7：财务优化&主力介入
    if stg_name == '小市值_基本面优化主力介入':
         # 计算ROE百分比排名，去除ROE较差的20%的股票
        df['ROE排名'] = df.groupby('交易日期')['ROE'].rank(ascending=False, method='min', pct=True)
        df = df[df['R_np_atoopc@xbx'] > 0]


        df['营业收入同比增速排名'] = df.groupby('交易日期')['营业收入同比增速_60'].rank(ascending=False, method='min', pct=True)
        condition = (df['营业收入同比增速排名'] < 0.6)
        # # 计算归母净利润同比增速排名
        df['归母净利润同比增速排名'] = df.groupby('交易日期')['归母净利润同比增速_60'].rank(ascending=False, method='min', pct=True)
        condition &= (df['归母净利润同比增速排名'] < 0.8)
        df['营业收入同比增速排名1'] = df.groupby('交易日期')['营业收入同比增速_60'].rank(ascending=False, method='min')


        # 计算Ret20排名  衡量超跌
        df['Ret_20排名'] = df.groupby('交易日期')['Ret_20'].rank(ascending=True, method='min')
        # 计算总市值排名
        df['总市值排名'] = df.groupby('交易日期')['总市值'].rank(ascending=True, method='min')
        # 中户资金
        df['中户买入占比排名'] = df.groupby('交易日期')['中户买入占比'].rank(ascending=False, method='min')
        # 计算成交额
        df['成交额std_20_排名'] = df.groupby('交易日期')['成交额std_20'].rank(ascending=True, method='min')


        df = df[df['上市至今交易天数'] > 250]
        df = df[df['总市值'] < 3e9]
        df = df[df['换手率_20'] > 0.01]
        df = df[df['毛利率'] > 0]
        df = df[condition]
        # 计算复合因子
        df['复合因子'] = df['总市值排名'] + df['中户买入占比排名'] + df['成交额std_20_排名'] + df['Ret_20排名'] + df['营业收入同比增速排名1']

    # 优化方向8：AH溢价率选股
    if stg_name == 'AH溢价率选股':
        df = df[df['AH溢价率'].isna() == False]
        df['AH溢价率排名'] = df.groupby('交易日期')['AH溢价率'].rank(ascending=True, method='min')
        
        df['ROE排名'] = df.groupby('交易日期')['ROE'].rank(ascending=False, method='min', pct=True)
        df = df[df['R_np_atoopc@xbx'] > 0]
        df = df[df['收盘价'] < 50]
        #中户资金
        # df['机构资金净流入排名'] = df.groupby('交易日期')['机构资金净流入'].rank(ascending=False, method='min')

        # 计算复合因子
        df['复合因子'] = df['AH溢价率排名'] + df['ROE排名']

    # 优化方向9：Ret超跌反弹
    if stg_name == 'Ret超跌反弹':

        df['毛利率排名'] = df.groupby('交易日期')['毛利率'].rank(ascending=False, method='min', pct=True)
        df = df[df['毛利率排名'] <= 0.75]
        
        df = df[df['B_total_equity_atoopc@xbx'] > 0]
        df = df[df['B_total_equity_atoopc@xbx_60'] > 0]
        
        df = df[df['换手率_sum'] >= 0.05]
        df = df[df['换手率_sum'] <= 0.18]
        df = df[df['总市值'] <= 5e9]
        
        df["WR_10_排名"] = df.groupby("交易日期")['WR_10'].rank(ascending=True, method="min", pct=True)
        df = df[df["WR_10_排名"] >= 0.4]
        # df["WR_3_排名"] = df.groupby("交易日期")['WR_3'].rank(ascending=True, method="min", pct=True)
        # df = df[df["WR_3_排名"] <= 0.9]

        df['Ret_30_排名'] = df.groupby('交易日期')['Ret_30'].rank(ascending=True, method='min')
        df['Ret_3_排名'] = df.groupby('交易日期')['Ret_3'].rank(ascending=True, method='min')
        df['成交额std_5_排名'] = df.groupby('交易日期')['成交额std_5'].rank(ascending=True, method='min')
        
        df["STR_20排名"] = df.groupby("交易日期")["STR_20"].rank(ascending=True, method="min")
        df["总市值_排名"] = df.groupby("交易日期")["总市值"].rank(ascending=True, method="min")
        df['散户资金净流入_排名'] = df.groupby("交易日期")["散户资金净流入"].rank(ascending=True, method="min")
        df = df[df['换手率_5'] < 0.03]
        df['复合因子'] = 2*df["STR_20排名"] + df["总市值_排名"] + df['Ret_30_排名'] + df['Ret_3_排名'] + df['散户资金净流入_排名']

    # 对因子进行排名
    df['排名'] = df.groupby('交易日期')['复合因子'].rank()

    # 选取排名靠前的股票
    df = df[df['排名'] <= select_stock_num]

    return df
